#include "encode.h"
#include "nal.h"
#include "h264_set.h"
#include "cabac.h"
#include "encoder_logger.h"
#include "cabac_encode.h"

extern void marcoblock_analyse(h264_ctx_t *ctx);
extern void macroblock_encode(h264_ctx_t *ctx);

static inline int bs_pos( bs_t *s )
{
    return( 8 * (s->p - s->p_start) + (sizeof(s->cur_bits) * 8) - s->i_left );
}

static void nal_start(nal_t *nal, h264_ctx_t *ctx, int i_type, int i_ref_idc )
{
    nal->i_ref_idc        = i_ref_idc;
    nal->i_type           = i_type;
    nal->b_long_startcode = 1;

    nal->i_payload = 0;
    nal->p_payload = ctx->bs.p_start + bs_pos(&ctx->bs) / 8;
}

static int nal_end( nal_t *nal, h264_ctx_t *ctx)
{
    uint8_t *end = ctx->bs.p_start + bs_pos(&ctx->bs) / 8;
    nal->i_payload = end - nal->p_payload;
    return 0;
}

static inline void bs_rbsp_trailing( bs_t *s )
{
    bs_write( s, 1, 1 );
    bs_write( s, s->i_left & 7, 0  );
}


static void write_sps(h264_ctx_t *ctx)
{
    bs_t *bs = &ctx->bs;
    struct sps_param_t *sps = &ctx->sps;
    nal_t nal = {0};
    nal_start(&nal, ctx, NAL_SPS, NAL_PRIORITY_HIGHEST);
    bs_write(bs, 8, sps->profile_idc);

    // set constraint_set_flag
    bs_write(bs, 1, sps->b_constraint_set0_flag);
    bs_write(bs, 1, sps->b_constraint_set1_flag);
    bs_write(bs, 1, sps->b_constraint_set2_flag);
    bs_write(bs, 1, sps->b_constraint_set3_flag);

    bs_write(bs, 4, 0); // reserved

    bs_write(bs, 8, sps->level_idc);

    bs_write_ue(bs, sps->id);
    if (sps->profile_idc >= PROFILE_HIGH) {
        bs_write_ue(bs, sps->chroma_format_idc);
        bs_write_ue(bs, BIT_DEPTH - 8); // bit depth       // bit_depth_luma_minus8
        bs_write_ue(bs, BIT_DEPTH - 8); // bit depth       // bit_depth_chroma_minus8
    }

    bs_write(bs, 1, sps->b_qpprime_y_zero_transform_bypass);
    bs_write(bs, 1, sps->seq_scaling_matrix_present_flag);  // 缩放矩阵

    bs_write_ue(bs, sps->i_log2_max_frame_num - 4);
    bs_write_ue(bs, sps->pic_order_cnt_type);
    bs_write_ue(bs, sps->i_num_ref_frames);
    bs_write(bs, 1, sps->b_gaps_in_frame_num_value_allowed);
    bs_write_ue(bs, sps->i_mb_width - 1);
    if (sps->b_frame_mbs_only == 1) {    // 帧模式
        bs_write_ue(bs, sps->i_mb_height - 1);
    } else {    // 场模式
        bs_write_ue(bs, (sps->i_mb_height >> 1) - 1);
    }

    bs_write(bs, 1, sps->b_frame_mbs_only);
    bs_write(bs, 1, sps->b_direct8x8_inference);
    bs_write(bs, 1, sps->b_crop);

    if (sps->b_crop == 1) {
        bs_write_ue(bs, sps->crop_left >> 1);
        bs_write_ue(bs, sps->crop_right >> 1);
        bs_write_ue(bs, sps->crop_top >> 1);
        bs_write_ue(bs, sps->crop_bottom >> 1);
    }
    bs_write(bs, 1, sps->vui_parameters_present_flag);
    bs_rbsp_trailing(bs);
    bs_flush(bs);
    nal_end(&nal, ctx);
    io_mng_write_nal_pack(&nal, OUTPUT_FILENAME);
    return;
}

void print_buff(uint8_t *buff, int size)
{
    for (int i = 0; i < size; ++i) {
        printf("%02x ", buff[i]);
    }
    printf("\n");
}

void write_pps(h264_ctx_t *ctx)
{
    bs_t *bs = &ctx->bs;
    struct pps_param_t *pps = &ctx->pps;
    nal_t nal = {0};
    nal_start(&nal, ctx, NAL_PPS, NAL_PRIORITY_HIGHEST);

    bs_write_ue(bs, pps->i_id);
    bs_write_ue(bs, pps->i_sps_id);
    bs_write(bs, 1, pps->b_cabac);
    bs_write(bs, 1, pps->b_pic_order);
    bs_write_ue(bs, pps->i_num_slice_groups - 1);
    bs_write_ue(bs, pps->i_num_ref_idx_l0_default_active - 1);
    bs_write_ue(bs, pps->i_num_ref_idx_l1_default_active - 1);
    bs_write(bs, 1, pps->b_weighted_pred);
    bs_write(bs, 2, pps->b_weighted_bipred);
    bs_write_se(bs, pps->i_pic_init_qp - 26);
    bs_write_se(bs, pps->i_pic_init_qs - 26);
    bs_write_se(bs, pps->i_chroma_qp_index_offset);

    bs_write(bs, 1, pps->b_deblocking_filter_control);
    bs_write(bs, 1, pps->b_constrained_intra_pred);
    bs_write(bs, 1, pps->b_redundant_pic_cnt);

    if (pps->b_transform_8x8_mode == 1) {
        bs_write(bs, 1, pps->b_transform_8x8_mode);
        bs_write(bs, 1, 0);
        bs_write_se(bs, pps->i_chroma_qp_index_offset);
    }

    bs_rbsp_trailing(bs);
    bs_flush(bs);
    nal_end(&nal, ctx);
    io_mng_write_nal_pack(&nal, OUTPUT_FILENAME);
    return;
}

static void slice_header_init(h264_ctx_t *ctx, h264_slice_header_t *slice_head)
{
    slice_head->i_first_mb = 0;
    slice_head->i_last_mb = ctx->cur_mb.i_mb_count - 1;
    slice_head->slice_type = SLICE_TYPE_I;     // 目前只编IDR slice
    slice_head->sps = &ctx->sps;
    slice_head->pps = &ctx->pps;
    slice_head->i_idr_pic_id = 0;
    slice_head->i_qp_delta = 5;
    slice_head->i_disable_deblocking_filter_idc = 0;
    slice_head->i_beta_offset = 0;
    slice_head->i_alpha_c0_offset = 0;
    slice_head->i_qp = 28;
    slice_head->i_cabac_init_idc = 0;

}

static void slice_header_write(h264_ctx_t *ctx)
{
    nal_start(&ctx->slice_nal, ctx, NAL_SLICE_IDR, NAL_PRIORITY_HIGHEST);
    bs_t *bs = &ctx->bs;
    h264_slice_header_t *sh = &ctx->slice_head;
    bs_write_ue(bs, sh->i_first_mb);
    bs_write_ue(bs, sh->slice_type + 5);    // FIXME: 不知道为什么+5
    bs_write_ue(bs, sh->pps->i_id);    // pic_parameter_set_id
    
    // frame_num
    bs_write(bs, sh->sps->i_log2_max_frame_num, sh->i_frame_num & ((1<<sh->sps->i_log2_max_frame_num)-1) );
    if( sh->i_idr_pic_id >= 0 ) /* NAL IDR */
    {
        bs_write_ue(bs, sh->i_idr_pic_id );
    }
    bs_write(bs, 1, 0 );  /* no output of prior pics flag */
    bs_write(bs, 1, 0 );  /* long term reference flag */
    bs_write_se(bs, sh->i_qp_delta);
    bs_write_ue(bs, sh->i_disable_deblocking_filter_idc);
    if (sh->i_disable_deblocking_filter_idc != 1) {
        bs_write_se(bs, sh->i_alpha_c0_offset);
        bs_write_se(bs, sh->i_beta_offset);
    }

}

static void init_pic_data_buffer(h264_ctx_t *ctx)
{
    // Y分量
    ctx->cur_mb.pic.p_fenc[0] = (pixel *)ctx->cur_mb.pic.fenc_buf;
    ctx->cur_mb.pic.p_fdec[0] = (pixel *)ctx->cur_mb.pic.fdec_buf + 2 * FDEC_STRIDE;

    // U分量
    ctx->cur_mb.pic.p_fenc[1] = (pixel *)ctx->cur_mb.pic.fenc_buf + 16 * FENC_STRIDE;
    ctx->cur_mb.pic.p_fdec[1] = (pixel *)ctx->cur_mb.pic.fdec_buf + 20 * FDEC_STRIDE;

    // V分量
    ctx->cur_mb.pic.p_fenc[2] = (pixel *)ctx->cur_mb.pic.fenc_buf + 16 * FENC_STRIDE + 8;
    ctx->cur_mb.pic.p_fdec[2] = (pixel *)ctx->cur_mb.pic.fdec_buf + 20 * FDEC_STRIDE + 16;
}

static void slice_write(h264_ctx_t *ctx)
{
    slice_header_write(ctx);
    if (ctx->pps.b_cabac == 1) {
        INFO_LOG("ctx->bs start %p pos %p offset %ld\n", ctx->bs.p_start, ctx->bs.p, 
            ctx->bs.p - ctx->bs.p_start);
        bs_align_1(&ctx->bs);
        h264_cabac_context_init(&ctx->cabac, SLICE_TYPE_I, ctx->slice_head.i_qp, ctx->slice_head.i_cabac_init_idc);
        h264_cabac_encode_init(&ctx->cabac, ctx->bs.p, ctx->bs.p_end);
    }

    ctx->cur_mb.i_last_qp = ctx->slice_head.i_qp;
    ctx->cur_mb.i_last_dqp = 0;
    ctx->cur_mb.field_decoding_flag = 0;
    ctx->cur_mb.i_qp = 28;
    int i_mb_x = 0;
    int i_mb_y = 0;
    int mb_xy = 0;
    init_pic_data_buffer(ctx);

    // 开始编码
    while (1)
    {
        mb_xy = i_mb_x + i_mb_y * ctx->sps.i_mb_width;
        INFO_LOG("start encode slice. cur mb_xy %d mb_x %d mb_y %d\n", mb_xy, i_mb_x, i_mb_y);
        if (i_mb_x == 0) {
            // TODO:filter
        }

        macroblock_cache_load(ctx, i_mb_x, i_mb_y);
        marcoblock_analyse(ctx);
        INFO_LOG("16x16 predict %d chroma predict %d\n", ctx->cur_mb.i_intra16x16_pred_mode, 
                ctx->cur_mb.i_chroma_pred_mode);
        macroblock_encode(ctx);
        cabac_encode(ctx, &ctx->cabac);
        // ctx->
        // TEST:
        break;
        
        if( mb_xy == ctx->slice_head.i_last_mb ) { // 编码的宏块数到达最后一个宏块，退出
            break;
        }

        i_mb_x++;
        if( i_mb_x == ctx->cur_mb.i_mb_width )
        {
            i_mb_y++;
            i_mb_x = 0;
        }
    }
    
    bs_rbsp_trailing(&ctx->bs);
    bs_flush(&ctx->bs);
    nal_end(&ctx->slice_nal, ctx);
    io_mng_write_nal_pack(&ctx->slice_nal, OUTPUT_FILENAME);
}

void h264_encode(h264_ctx_t *ctx)
{
    write_sps(ctx);
    write_pps(ctx);
    slice_header_init(ctx, &ctx->slice_head);
    slice_write(ctx);
    return;
}